package com.starhub.integration.sdk.custom.core;

import lombok.Builder;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.starhub.utils.Util;

import java.io.IOException;
import java.util.List;
import java.util.Map;

/**
 * 私有化向量模型客户端
 * 支持基于xinfrence部署的bge-m3等向量模型
 * @version: V1.0
 * @author: liuhf
 * @time 2025-08-08
 */
@Builder
public class CustomVectorClient {
    private OkHttpClient okHttpClient;
    private String apiHost;
    private String modelName;
    private String apiPath;

    /**
     * 获取私有化向量模型的embedding
     * 
     * @param text 需要向量化的文本
     * @return 文本的向量表示
     * @throws IOException 当API调用失败或响应解析失败时抛出
     */
    public float[] getEmbedding(String text) throws IOException {
        if (Util.isEmpty(text)) {
            throw new IOException("输入文本不能为空");
        }

        String url = apiHost + apiPath;

        // 构建请求参数
        Map<String, Object> requestParams = Map.of(
            "model", modelName,
            "input", List.of(text)
        );

        RequestBody body = RequestBody.create(
                JSON.toJSONString(requestParams),
                MediaType.parse("application/json; charset=utf-8")
        );

        Request.Builder requestBuilder = new Request.Builder()
                .url(url)
                .addHeader("Content-Type", "application/json")
                .post(body);

        // 如果配置了API Key，添加到请求头
        /*if (Util.isNotEmpty(token)) {
            requestBuilder.addHeader("Authorization", "Bearer " + token);
        }*/

        Request request = requestBuilder.build();

        try (Response response = okHttpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException(String.format("API请求失败，状态码: %d", response.code()));
            }

            String responseBody = response.body().string();
            if (Util.isEmpty(responseBody)) {
                throw new IOException("API返回响应为空");
            }

            return parseEmbeddingResponse(responseBody);
        } catch (IOException e) {
            throw new IOException("调用私有化向量API失败: " + e.getMessage(), e);
        }
    }

    /**
     * 解析私有化向量模型的embedding响应
     * 支持多种响应格式，包括OpenAI兼容格式和自定义格式
     * 
     * @param responseBody API响应内容
     * @return 向量数组
     * @throws IOException 当响应解析失败时抛出
     */
    private float[] parseEmbeddingResponse(String responseBody) throws IOException {
        try {
            JSONObject jsonObject = JSON.parseObject(responseBody);

            // 检查错误信息
            if (jsonObject.containsKey("error")) {
                String errorMsg = jsonObject.getString("error");
                throw new IOException("API返回错误: " + errorMsg);
            }

            // 尝试解析OpenAI兼容格式
            if (jsonObject.containsKey("data")) {
                JSONArray dataArray = jsonObject.getJSONArray("data");
                if (dataArray != null && !dataArray.isEmpty()) {
                    JSONObject dataObject = dataArray.getJSONObject(0);
                    JSONArray embeddingArray = dataObject.getJSONArray("embedding");
                    if (embeddingArray != null && !embeddingArray.isEmpty()) {
                        return convertToFloatArray(embeddingArray);
                    }
                }
            }

            // 尝试解析自定义格式（直接返回向量数组）
            if (jsonObject.containsKey("embeddings")) {
                JSONArray embeddingsArray = jsonObject.getJSONArray("embeddings");
                if (embeddingsArray != null && !embeddingsArray.isEmpty()) {
                    JSONArray embeddingArray = embeddingsArray.getJSONArray(0);
                    if (embeddingArray != null && !embeddingArray.isEmpty()) {
                        return convertToFloatArray(embeddingArray);
                    }
                }
            }

            // 尝试解析直接返回向量数组的格式
            if (jsonObject.containsKey("embedding")) {
                JSONArray embeddingArray = jsonObject.getJSONArray("embedding");
                if (embeddingArray != null && !embeddingArray.isEmpty()) {
                    return convertToFloatArray(embeddingArray);
                }
            }

            throw new IOException("无法解析API响应格式，响应内容: " + responseBody);
        } catch (Exception e) {
            throw new IOException("解析API响应失败: " + e.getMessage(), e);
        }
    }

    /**
     * 将JSONArray转换为float数组
     * 
     * @param array JSON数组
     * @return float数组
     */
    private float[] convertToFloatArray(JSONArray array) {
        float[] result = new float[array.size()];
        for (int i = 0; i < array.size(); i++) {
            result[i] = Util.getFloatValue(array.get(i), 0.0f);
        }
        return result;
    }
} 